import cv2
import numpy as np
import hashlib
import qrcode
from PIL import Image
from watermarking.attack import Attack
import imagehash
import matplotlib.pyplot as plt
import os
from collections import OrderedDict
import shutil

ACTION_MAP = {
    "blur": Attack.blur,
    "rotate180": Attack.rotate180,
    "rotate90": Attack.rotate90,
    "chop5": Attack.chop5,
    "chop10": Attack.chop10,
    "chop30": Attack.chop30,
    "gray": Attack.gray,
    "saltnoise": Attack.saltnoise,
    "randline": Attack.randline,
    "cover": Attack.cover,
    "brighter10": Attack.brighter10,
    "darker10": Attack.darker10,
    "largersize": Attack.largersize,
    "smallersize": Attack.smallersize,
}

# img_jump = [
#     "ILSVRC2012_val_00000037",
#     "ILSVRC2012_val_00005706",
#     "ILSVRC2012_val_00001531",
#     "ILSVRC2012_val_00001968",
#     "ILSVRC2012_val_00000147",
#     "ILSVRC2012_val_00001637",
#     "ILSVRC2012_val_00001060",
#     "ILSVRC2012_val_00000481",
#     "ILSVRC2012_val_00001731",
#     "ILSVRC2012_val_00000562",
#     "ILSVRC2012_val_00001689",
#     "ILSVRC2012_val_00001730",
#     "ILSVRC2012_val_00000123",
#     "ILSVRC2012_val_00002206",
#     "ILSVRC2012_val_00001362",
#     "ILSVRC2012_val_00001120",
#     "ILSVRC2012_val_00000145",
#     "ILSVRC2012_val_00000139",
#     "ILSVRC2012_val_00001078",
#     "ILSVRC2012_val_00000691",
#     "ILSVRC2012_val_00000620",
#     "ILSVRC2012_val_00003145",
#     "ILSVRC2012_val_00032029",
#     "ILSVRC2012_val_00001805",
#     "ILSVRC2012_val_00000750",
#     "ILSVRC2012_val_00004724",
#     "ILSVRC2012_val_00001036",
#     "ILSVRC2012_val_00002570",
#     "ILSVRC2012_val_00004530",
#     "ILSVRC2012_val_00007422",
#     "ILSVRC2012_val_00003074",
#     "ILSVRC2012_val_00002003",
#     "ILSVRC2012_val_00000794",
#     "ILSVRC2012_val_00002953",
#     "ILSVRC2012_val_00000878",
#     "ILSVRC2012_val_00000774",
#     "ILSVRC2012_val_00000735",
#     "ILSVRC2012_val_00001371",
#     "ILSVRC2012_val_00001737",
#     "ILSVRC2012_val_00002076",
#     "ILSVRC2012_val_00003163",
#     "ILSVRC2012_val_00003125",
#     "ILSVRC2012_val_00003314",
#     "ILSVRC2012_val_00001027",
#     "ILSVRC2012_val_00000892",
#     "ILSVRC2012_val_00040420",
#     "ILSVRC2012_val_00039996",
#     "ILSVRC2012_val_00049287",
#     "ILSVRC2012_val_00030656",
# ]

img_jump = []

att_names = [k for k in ACTION_MAP.keys()]
test_result_dir = "/mnt/e/pydev/hash-store/watermarking/images/"


def PIL_to_cv2(img, mode=None):
    if mode is None:
        # Gray Mode
        return np.array(img.getdata()).reshape(img.size).astype(np.float32)
        # RGB to opencv BGR
    else:
        return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


def cv2_to_PIL(img):
    cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)


def make_qrcode(qr_data):
    qr = qrcode.QRCode(
        # version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_H,
        box_size=10,
        border=4,
    )
    qr.add_data(qr_data)
    wm = qr.make_image(fill_color="black", back_color="white")
    logo = Image.open("/mnt/e/pydev/hash-store/logo.png")
    logo = logo.resize((int(logo.size[0] * 0.5), int(logo.size[1] * 0.5)))
    logo_position = ((wm.size[0] - logo.size[0]) // 2, (wm.size[1] - logo.size[1]) // 2)
    wm.paste(logo, logo_position)
    # wm.paste(location, (0, 0))
    wm = np.array(wm).reshape(wm.size).astype(np.float32)
    return wm


def make_qrcode_test(imraw_file):
    pic = open(imraw_file, "rb")
    img = cv2.imread(imraw_file)

    # Get Qrcode information
    hashfunc = hashlib.sha256()
    hashfunc.update(pic.read())
    hashv = hashfunc.hexdigest()
    qr = qrcode.QRCode(
        # version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_H,
        box_size=10,
        border=4,
    )
    qr.add_data(hashv)
    # Generate image watermark
    wm = qr.make_image(fill_color="black", back_color="white")

    # location = np.zeros((10, 10), dtype=np.float32)
    # location = Image.fromarray(location)
    logo = Image.open("/mnt/e/pydev/hash-store/logo.png")
    logo = logo.resize((int(logo.size[0] * 0.5), int(logo.size[1] * 0.5)))
    logo_position = ((wm.size[0] - logo.size[0]) // 2, (wm.size[1] - logo.size[1]) // 2)
    wm.paste(logo, logo_position)
    # wm.paste(location, (0, 0))
    wm = PIL_to_cv2(wm)
    return img, wm


# 14 kinds attacks to wmd image
# return attack type, attackd images, and distance to wmd image.
def attack_analysis(wmd):
    phash_wmd = imagehash.phash(Image.fromarray(wmd),hash_size=11)
    chash_wmd = imagehash.crop_resistant_hash(Image.fromarray(wmd), imagehash.phash)
    # print(chash_wmd)
    
    att_hashs = []
    att_chashs = []
    att_imgs = []
    att_names = []
    for k in ACTION_MAP.keys():
        att_names.append(k)
        att_img = ACTION_MAP[k](wmd)
        att_hash = imagehash.phash(Image.fromarray(att_img),11)
        att_chash = imagehash.crop_resistant_hash(
            Image.fromarray(att_img), imagehash.phash
        )

        att_hashs.append(att_hash)
        att_chashs.append(att_chash)
        att_imgs.append(att_img)

    distances = [ah - phash_wmd for ah in att_hashs]
    cdistances = [ah - chash_wmd for ah in att_chashs]
    return att_names, att_imgs, distances, cdistances


def attack_imgs_plt(att_imgs):
    rows = 2
    cols = 7
    axes = []
    fig = plt.figure()
    for a in range(rows * cols):
        # b = np.random.randint(7, size=(height, width))
        axes.append(fig.add_subplot(rows, cols, a + 1))
        axes[-1].set_title(att_names[a])
        plt.imshow(att_imgs[a])
        plt.xticks([])
        plt.yticks([])
    # fig.tight_layout()
    plt.subplots_adjust(hspace=0.1, wspace=0.1)
    plt.show()


def attack_imgs_plt_from_file(attacked_dir):
    rows = 2
    cols = 7
    axes = []

    # attacked_dir = "/mnt/e/pydev/hash-store/imgs/attacked"
    imgs_file = os.listdir(attacked_dir)
    imgs = []
    for imf in imgs_file:
        bgr_img = cv2.imread(f"{attacked_dir}/{imf}")
        rgb_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)
        imgs.append(rgb_img)

    fig = plt.figure(figsize=(25, 8))
    for a in range(rows * cols):
        # b = np.random.randint(7, size=(height, width))
        axes.append(fig.add_subplot(rows, cols, a + 1))
        # axes[-1].set_title(att_names[a])
        plt.imshow(imgs[a])
        plt.xticks([])
        plt.yticks([])
    # fig.tight_layout()
    plt.subplots_adjust(
        hspace=-0.5, wspace=0.05, left=0.01, right=0.99, top=1.2, bottom=-0.17
    )
    plt.show()


def a_clazz_img_wmd_plt(imgs):
    rows = 5
    cols = 8
    axes = []
    fig = plt.figure()
    # fig.tight_layout(h_pad=0)
    for a in range(rows * cols):
        # b = np.random.randint(7, size=(height, width))
        axes.append(fig.add_subplot(rows, cols, a + 1))
        # axes[-1].set_title(att_names[a])
        rgb_img = cv2.cvtColor(imgs[a], cv2.COLOR_BGR2RGB)
        plt.imshow(rgb_img)
        plt.xticks([])
        plt.yticks([])
    plt.subplots_adjust(hspace=0.1, wspace=0.1)

    plt.show()
    pass


def a_clazz_distance_heat_plt(arr):
    # plt.colorbar()
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            text = plt.text(
                j, i, arr[i, j], ha="center", va="center", color="w", size=12
            )
    plt.imshow(arr)
    plt.tight_layout()
    plt.show()


def different_clazz_plt():
    distances = [
        0.7157895,
        31.915789,
        31.810526,
        6.8,
        13.663157,
        27.621052,
        2.4421053,
        0.44210526,
        4.3157897,
        11.28421,
        11.452632,
        0.21052632,
        0.12631579,
        0.16842106,
    ]
    cdistances = [
        0.10060679,
        4.03125,
        4.05,
        0.52983724,
        0.91043527,
        2.73162667,
        1.06969587,
        0.06259611,
        4.07711589,
        1.01912698,
        2.13563368,
        1.38891555,
        0.00767029,
        0.01365857,
    ]

    distances = np.array(distances)
    labels = np.array(att_names)
    d_dict = OrderedDict()

    for d in range(len(distances)):
        d_dict[labels[d]] = (distances[d], cdistances[d])

    d_dict = OrderedDict(sorted(d_dict.items(), key=lambda t: t[1][0]))
    distances = [v[0] for v in d_dict.values()]
    cdistances = [v[1] for v in d_dict.values()]

    edgecolors = ["#8FBC94"] * 14
    # edgecolors.extend(["#8FBC94"] * 6)
    edgecolors1 = ["#EFDC05"] * 14

    fig, ax = plt.subplots(figsize=(16, 8))
    # ax.bar(fruits, counts, label=bar_labels, color=bar_colors)

    ax.set_ylabel("Hamming Distance", size=16)
    # ax.set_title("Fruit supply by kind and color")
    # ax.legend(title="Fruit color")
    ax.axhline(y=10, color="gray", linestyle="dotted")

    x = np.arange(len(d_dict))

    bar2 = ax.bar(
        x - 0.2,
        list(cdistances),
        # hatch="\\",
        # color="w",
        color=edgecolors1,
        edgecolor="#30A9DE",
        width=0.4,
        label="Crop Resistant Hash",
    )

    bar1 = ax.bar(
        x + 0.2,
        list(distances),
        # hatch="/",
        # color="w",
        color=edgecolors,
        edgecolor="#30A9DE",
        width=0.4,
        label="Perceptual Hash",
    )

    plt.xticks(rotation=30, fontsize=16)
    plt.yticks(fontsize=16)
    ax.bar_label(bar1, label_type="edge", fmt="%.2f", size=12)
    ax.bar_label(bar2, label_type="edge", fmt="%.2f", size=12)
    ax.set_xticks(x, list(d_dict.keys()))
    plt.subplots_adjust(left=0.05, right=0.99, top=0.99, bottom=0.12)
    plt.legend(fontsize=16)
    plt.show()


def different_clazz_qr_plt():
    qr_rate_dict = {
        "blur": [68, 3],
        "rotate180": [68, 9],
        "rotate90": [68, 2],
        "chop5": [68, 36],
        "chop10": [68, 3],
        "chop30": [68, 3],
        "gray": [68, 67],
        "saltnoise": [68, 42],
        "randline": [68, 64],
        "cover": [68, 65],
        "brighter10": [68, 35],
        "darker10": [68, 67],
        "largersize": [68, 4],
        "smallersize": [68, 4],
    }

    for qrd in qr_rate_dict.keys():
        qr_rate_dict[qrd] = qr_rate_dict[qrd][1] / qr_rate_dict[qrd][0]

    print(qr_rate_dict)

    d_dict = OrderedDict()
    d_dict = OrderedDict(sorted(qr_rate_dict.items(), key=lambda t: t[1]))

    print(d_dict)
    fig, ax = plt.subplots(figsize=(16, 8))
    # ax.bar(fruits, counts, label=bar_labels, color=bar_colors)

    ax.set_ylabel("Recognition Rate", size=16)
    # ax.set_title("Fruit supply by kind and color")
    # ax.legend(title="Fruit color")
    ax.axhline(y=0.5, color="gray", linestyle="dotted")
    edgecolors = ["#EFDC05"] * 7
    edgecolors.extend(["#8FBC94"] * 7)
    bar1 = ax.bar(
        list(d_dict.keys())[:7],
        list(d_dict.values())[:7],
        # hatch="/",
        # color="w",
        color=edgecolors[:7],
        edgecolor="#30A9DE",
        width=0.5,
        label="Can resist attacks",
    )
    ax.bar_label(bar1, label_type="edge", fmt="%.2f", size=16)

    bar1 = ax.bar(
        list(d_dict.keys())[7:],
        list(d_dict.values())[7:],
        # hatch="/",
        # color="w",
        color=edgecolors[7:],
        edgecolor="#30A9DE",
        width=0.5,
        label="Can not resist attacks",
    )
    ax.bar_label(bar1, label_type="edge", fmt="%.2f", size=16)

    plt.subplots_adjust(left=0.05, right=0.99, top=0.99, bottom=0.12)
    plt.xticks(rotation=30, fontsize=16)
    plt.yticks(fontsize=16)
    plt.legend(fontsize=16)
    plt.show()


def hash_process_all():
    process_all_path = "/mnt/e/pydev/hash-store/imgs/processall"
    imgs = os.listdir(process_all_path)
    for img in imgs:
        img_file = f"{process_all_path}/{img}"
        phash = imagehash.phash(Image.open(img_file))
        with open(img_file, "rb") as img_f:
            hashf = hashlib.sha256()
            hashf.update(img_f.read())
            chash = hashf.hexdigest()

        print(img, phash, chash)
